import requests
import unittest
import json
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, Any

class TestAgentAPI(unittest.TestCase):
    """Test cases for Agent Service API."""
    
    def setUp(self):
        """Set up test environment."""
        self.base_url = "http://localhost:8000"
        self.headers = {"Content-Type": "application/json"}
        self.test_task_id = None

    def test_01_list_tools(self):
        """Test GET /agent/tools endpoint."""
        response = requests.get(f"{self.base_url}/agent/tools")
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIn("tools", data)
        self.assertIsInstance(data["tools"], list)

    def test_02_create_task(self):
        """Test POST /agent/task endpoint."""
        payload = {
            "description": "Generate a presentation about AI",
            "parameters": {"topic": "AI"}
        }
        response = requests.post(
            f"{self.base_url}/agent/task",
            headers=self.headers,
            json=payload
        )
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIn("task_id", data)
        self.assertIn("status", data)
        self.test_task_id = data["task_id"]

    def test_03_get_task_status(self):
        """Test GET /agent/task/{task_id} endpoint."""
        if not self.test_task_id:
            self.test_02_create_task()
        
        response = requests.get(f"{self.base_url}/agent/task/{self.test_task_id}")
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertEqual(data["task_id"], self.test_task_id)

    def test_04_process_task(self):
        """Test POST /agent/task/{task_id}/process endpoint."""
        if not self.test_task_id:
            self.test_02_create_task()
        
        response = requests.post(f"{self.base_url}/agent/task/{self.test_task_id}/process")
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIn("status", data)

    def test_05_generate_ppt(self):
        """Test POST /agent/generate-ppt endpoint."""
        payload = {
            "topic": "Artificial Intelligence Trends 2024",
            "output_filename": "test_presentation.pptx",
            "parameters": {
                "num_slides": 5,
                "include_outline": True,
                "style": "professional"
            }
        }
        response = requests.post(
            f"{self.base_url}/agent/generate-ppt",
            headers=self.headers,
            json=payload
        )
        self.assertEqual(response.status_code, 200)
        data = response.json()
        self.assertIn("status", data)
        self.assertIn("file_path", data)

    def test_06_download_ppt(self):
        """Test GET /agent/download-ppt/{filename} endpoint."""
        # First generate a PPT
        self.test_05_generate_ppt()
        filename = "test_presentation.pptx"
        
        response = requests.get(f"{self.base_url}/agent/download-ppt/{filename}")
        self.assertEqual(response.status_code, 200)
        self.assertIn("application/vnd.openxmlformats-officedocument.presentationml.presentation",
                     response.headers["Content-Type"])

    def test_07_error_handling(self):
        """Test error handling scenarios."""
        # Test invalid task ID
        response = requests.get(f"{self.base_url}/agent/task/invalid-id")
        self.assertEqual(response.status_code, 404)

        # Test invalid JSON
        response = requests.post(
            f"{self.base_url}/agent/task",
            headers=self.headers,
            data="invalid json"
        )
        self.assertEqual(response.status_code, 422)

        # Test missing required fields
        response = requests.post(
            f"{self.base_url}/agent/generate-ppt",
            headers=self.headers,
            json={}
        )
        self.assertEqual(response.status_code, 422)

    def test_08_concurrent_requests(self):
        """Test handling of concurrent requests."""
        def create_task(i: int) -> Dict[str, Any]:
            payload = {
                "description": f"Concurrent task {i}",
                "parameters": {"priority": i}
            }
            response = requests.post(
                f"{self.base_url}/agent/task",
                headers=self.headers,
                json=payload
            )
            return response.json()

        with ThreadPoolExecutor(max_workers=5) as executor:
            futures = [executor.submit(create_task, i) for i in range(5)]
            results = [f.result() for f in futures]

        # Verify all tasks were created successfully
        self.assertEqual(len(results), 5)
        for result in results:
            self.assertIn("task_id", result)
            self.assertIn("status", result)

    def test_09_task_workflow(self):
        """Test complete task workflow."""
        # 1. Create task
        create_response = requests.post(
            f"{self.base_url}/agent/task",
            headers=self.headers,
            json={
                "description": "Create a comprehensive presentation about Python",
                "parameters": {
                    "topic": "Python Programming",
                    "subtopics": ["Basics", "OOP", "Advanced Features"]
                }
            }
        )
        self.assertEqual(create_response.status_code, 200)
        task_data = create_response.json()
        task_id = task_data["task_id"]

        # 2. Start processing
        process_response = requests.post(f"{self.base_url}/agent/task/{task_id}/process")
        self.assertEqual(process_response.status_code, 200)

        # 3. Poll for completion (with timeout)
        max_attempts = 10
        for _ in range(max_attempts):
            status_response = requests.get(f"{self.base_url}/agent/task/{task_id}")
            self.assertEqual(status_response.status_code, 200)
            status_data = status_response.json()
            
            if status_data["status"] in ["completed", "failed"]:
                break
                
            time.sleep(1)  # Wait before next poll

        self.assertIn(status_data["status"], ["completed", "failed"])

if __name__ == '__main__':
    unittest.main(verbosity=2) 